package com.guchenbo.example.rsa;

import org.apache.commons.codec.binary.Base64;

import javax.crypto.Cipher;
import java.security.*;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.HashMap;
import java.util.Map;

import static org.apache.commons.codec.binary.Base64.decodeBase64;

/**
 * @Author gucb
 * @Date 2017/5/11.
 */
public class RSAUtils {
    public final static String ALGORITHM = "RSA";
    private static Map<String, String> keyMap = new HashMap<>();

    private static String PUBLIC_KEY = "public_key";
    private static String PRIVATE_KEY = "private_key";

    public static String encryptByPublic(String data, String publicKey) {
        try {
            KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM);
            X509EncodedKeySpec keySpec = new X509EncodedKeySpec(keyByte(publicKey));
            PublicKey pubKey = keyFactory.generatePublic(keySpec);

            Cipher cipher = Cipher.getInstance(keyFactory.getAlgorithm());
            cipher.init(Cipher.ENCRYPT_MODE, pubKey);

            byte[] encryptedData = cipher.doFinal(data.getBytes("utf-8"));
            return Base64.encodeBase64String(encryptedData);
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    public static String decryptByPrivate(String data, String privateKey) {
        try {
            KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM);
            PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(keyByte(privateKey));
            PrivateKey priKey = keyFactory.generatePrivate(keySpec);

            Cipher cipher = Cipher.getInstance(keyFactory.getAlgorithm());
            cipher.init(Cipher.DECRYPT_MODE, priKey);

            byte[] bytes = Base64.decodeBase64(data);
            byte[] ret = cipher.doFinal(bytes);

            return new String(ret, "utf-8");
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    public static void initKeyPair() {
        try {
            KeyPairGenerator generator = KeyPairGenerator.getInstance(ALGORITHM);
            generator.initialize(1024);
            KeyPair pair = generator.generateKeyPair();

            PublicKey pubKey = pair.getPublic();
            PrivateKey priKey = pair.getPrivate();

            byte[] pk = pubKey.getEncoded();
            byte[] pk2 = priKey.getEncoded();

            keyMap.put(PUBLIC_KEY, Base64.encodeBase64String(pk));
            keyMap.put(PRIVATE_KEY, Base64.encodeBase64String(pk2));
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public static String getPublicKey() {
        return keyMap.get(PUBLIC_KEY);
    }

    public static String getPrivateKey() {
        return keyMap.get(PRIVATE_KEY);
    }

    private static byte[] keyByte(String key) {
        return decodeBase64(key.getBytes());
    }
}
